import torch


# 修改预训练权重成适合自己的分类
def change_feature(check_point, num_classes):
    """
    修改预训练权重中的参数
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    check_point = torch.load(check_point, map_location=device)

    import collections
    dicts = collections.OrderedDict()

    for k, value in check_point.items():
        print("names:{}".format(k))
        print("shape:{}".format(value.size()))     # 打印出权重文件中的节点名称和形状

        if k == "decoder.embedding.weight":        # 根据打印出网络层，修改对应名称对应的值
            value = torch.ones(num_classes, 256)
        if k == "decoder.out.weight":
            value = torch.ones(num_classes, 256)
        if k == "decoder.out.bias":
            value = torch.ones(num_classes)
        dicts[k] = value

    torch.save(dicts, "model/changWeight/change_weight.pth")


if __name__ == "__main__":
    model_path = "model/pretrained/decoder.pth"
    num_classes = 77     # 自己分类的类别数
    change_feature(model_path, 77)